"""
---------------------------------------------------------------------
- File: middleware.py
- Date: 2025/11/14
- Author: yangzhide
- Email: xxzjyzd@126.com
- Description: 中间件
---------------------------------------------------------------------
"""

import json
import time
import uuid

from fastapi import Request, Response
from starlette.background import BackgroundTask
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.types import ASGIApp

from app.core.logs import logger
from app.db.session import AsyncSessionLocal
from app.models.system import Logs


def format_log_body(body_bytes: bytes) -> dict | list | str | None:
    """尝试将 bytes 解析为 JSON，否则返回字符串预览"""
    if not body_bytes:
        return None
    try:
        return json.loads(body_bytes.decode("utf-8"))
    except Exception:  # noqa
        return body_bytes.decode("utf-8", errors="ignore")  # 直接返回


async def log_to_database(log_data: dict):
    """异步将日志数据存入数据库"""
    async with AsyncSessionLocal() as session:
        try:
            log_entry = Logs(**log_data)
            session.add(log_entry)
            await session.commit()
        except Exception as e:
            await session.rollback()
            logger.error(f"日志记录失败（不影响正常的请求响应）: {e}")


class LogMiddleware(BaseHTTPMiddleware):
    def __init__(self, app: ASGIApp):
        super().__init__(app)

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        # ---请求阶段---
        request_id = uuid.uuid4()
        # request.state对象存储请求相关数据，可在请求生命周期内使用
        request.state.request_id = request_id

        start_time = time.time()
        req_body_bytes = await request.body()

        # ---放行阶段---
        response = await call_next(request)

        # ---响应阶段---
        end_time = time.time()
        process_time = (end_time - start_time) * 1000  # 毫秒

        # 获取用户信息
        auth_user = getattr(request.state, "auth_user", None)
        user_id = auth_user.id if auth_user else None

        # 读取响应体
        res_body_bytes = b""
        async for chunk in response.body_iterator:
            res_body_bytes += chunk

        content_type = response.headers.get("content-type", "")
        if "application/json" in content_type:
            response_body_to_log = format_log_body(res_body_bytes)
        else:
            # 对于非 JSON 响应 (如 HTML, text/plain)，我们只记录一个提示信息
            # 或者记录解码后的文本预览
            response_body_to_log = {
                "content_type": content_type,
                "content_preview": res_body_bytes.decode("utf-8", errors="ignore")[:500] + "...",
            }
        # 创建日志对象
        log_data = {
            "request_id": request_id,
            "ip": request.client.host if request.client else None,
            "method": request.method,
            "url": str(request.url),
            "path_params": dict(request.path_params) or None,
            "query_params": dict(request.query_params) or None,
            "request_body": format_log_body(req_body_bytes) if req_body_bytes else None,
            "status_code": response.status_code,
            "response_body": response_body_to_log,
            "process_time": process_time,
            "user_id": user_id,
        }

        # 创建BackgroundTask，异步将日志数据存入数据库
        # 在Uvicorn成功将响应发送完毕后，它会检查Response对象上是否有.background属性。
        # 如果存在，Uvicorn会await 执行这个属性中定义的所有任务。无论是API层添加的任务，
        # 还是中间件层添加的任务，都会在这个阶段被执行。 它们与客户端的连接已经没有关系了。
        # 客户端已经收到了它的响应，可能已经断开连接了。
        # 这些任务在服务器的事件循环中作为独立的协程运行，直到完成。
        task = BackgroundTask(log_to_database, log_data=log_data)

        # 添加请求ID到响应头，方便追踪
        response.headers["X-Request-ID"] = str(request_id)
        # 返回原始响应
        return Response(
            content=res_body_bytes,
            status_code=response.status_code,
            headers=dict(response.headers),
            media_type=response.media_type,
            background=task,  # 添加background属性， uvicorn会await异步日志记录任务
        )
